import numpy as np

from src.impl.output import get_output_float16, get_output_float32
from matplotlib import pyplot as plt
from src.generator_utils import get_random_seed_tensor


def get_mre_and_mare(tensor, operator, framework, epsilon= 1e-7):

    # 得到tf和pytorch输出的numpy数组
    tf_output_b, torch_output_b, mnn_output_b, variable = get_output_float32(tensor, operator)

    # 获得三组输出的差异
    tf_torch_abs = np.maximum(tf_output_b - torch_output_b, torch_output_b - tf_output_b)
    tf_mnn_abs = np.maximum(tf_output_b - mnn_output_b, mnn_output_b - tf_output_b)
    torch_mnn_abs = np.maximum(torch_output_b - mnn_output_b, mnn_output_b - torch_output_b)
    # 获得总和
    tf_torch_var = np.sum(tf_torch_abs)
    tf_mnn_var = np.sum(tf_mnn_abs)
    torch_mnn_var = np.sum(torch_mnn_abs)
    # 比较并得到benchmark
    if tf_torch_var <= tf_mnn_var and tf_torch_var <= torch_mnn_var:
        benchmark = tf_output_b / 2 + torch_output_b / 2
    elif tf_mnn_var <= tf_torch_var and tf_mnn_var <= torch_mnn_var:
        benchmark = tf_output_b / 2 + mnn_output_b / 2
    else:
        benchmark = torch_output_b / 2 + mnn_output_b / 2
    tf_output_16, torch_output_16, mnn_output_16, _ = get_output_float16(tensor, operator, variable)

    # 差异矩阵
    diff_matrix_tf = tf_output_16 - benchmark
    diff_matrix_torch = torch_output_16 - benchmark
    diff_matrix_mnn = mnn_output_16 - benchmark

    # 求和项(绝对值)
    # epsilon取一个很小的值防止得0
    single_item_matrix_abs_tf = np.maximum(diff_matrix_tf, -diff_matrix_tf) / \
                            (np.maximum(tf_output_16, -tf_output_16) + epsilon)
    single_item_matrix_abs_torch = np.maximum(diff_matrix_torch, -diff_matrix_torch) / \
                            (np.maximum(torch_output_16, -torch_output_16) + epsilon)
    single_item_matrix_abs_mnn = np.maximum(diff_matrix_mnn, -diff_matrix_mnn) / \
                            (np.maximum(mnn_output_16, -mnn_output_16) + epsilon)

    if framework == 'tf':
        return np.sum(single_item_matrix_abs_tf) / np.size(single_item_matrix_abs_tf), \
               np.max(single_item_matrix_abs_tf)
    elif framework == 'torch':
        return np.sum(single_item_matrix_abs_torch) / np.size(single_item_matrix_abs_torch), \
               np.max(single_item_matrix_abs_torch)
    elif framework == 'mnn':
        return np.sum(single_item_matrix_abs_mnn) / np.size(single_item_matrix_abs_mnn), \
               np.max(single_item_matrix_abs_mnn)
    else:
        mre = (np.sum(single_item_matrix_abs_tf) / np.size(single_item_matrix_abs_tf)
               + np.sum(single_item_matrix_abs_torch) / np.size(single_item_matrix_abs_torch)
               + np.sum(single_item_matrix_abs_mnn) / np.size(single_item_matrix_abs_mnn)) / 3
        mare = (np.max(single_item_matrix_abs_tf) + np.max(single_item_matrix_abs_torch)
                + np.max(single_item_matrix_abs_mnn)) / 3
        return mre, mare


if __name__ == '__main__':
    lst = []
    point_mre_x = []
    point_mare_y = []
    # for i in range(100):
    #     np_a = np.load('./data/tf/{0}.npy'.format(i))
    #     lst.append(np_a)
    # for i in range(100):
    #     np_a = np.load('./data/torch/{0}.npy'.format(i))
    #     lst.append(np_a)
    # for i in range(100):
    #     np_a = np.load('./data/mnn/{0}.npy'.format(i))
    #     lst.append(np_a)

    # for i in range(100):
    #     np_a = np.load('../data/tf/batch_normalization/{0}.npy'.format(i))
    #     lst.append(np_a)
    # for i in range(100):
    #     np_a = np.load('../data/torch/batch_normalization/{0}.npy'.format(i))
    #     lst.append(np_a)
    # for i in range(100):
    #     np_a = np.load('../data/torch/batch_normalization/{0}.npy'.format(i+2678))
    #     lst.append(np_a)
    for i in range(300):
        lst.append(np.load('./{}.npy'.format(i)))
    for i in range(300):
        tensor = lst[i]
        if i < 100:
            x, y = get_mre_and_mare(tensor, 'batch_normalization', 'tf')
        elif i < 200:
            x, y = get_mre_and_mare(tensor, 'batch_normalization', 'torch')
        else:
            x, y = get_mre_and_mare(tensor, 'batch_normalization', 'mnn')
        point_mre_x.append(x)
        point_mare_y.append(y)

    print(point_mre_x)
    print(point_mare_y)
    # fig = plt.figure()
    # fig.suptitle('MRE/MARE Distribution Generated By Predoo-ln1 On Batch_Norm')
    # fig.subplots_adjust(hspace=0.4)
    # ax1 = fig.add_subplot(221)
    # ax1.set_title('All Samples')
    # ax1.scatter(point_mre_x, point_mare_y, c=np.random.random(300), cmap='rainbow')
    # ax1.set_xlabel("MRE",fontsize=7)
    # ax1.set_ylabel("MARE",fontsize=7, labelpad=6)
    # ax1.set_xticklabels([0.00, 0.00, 0.25, 0.50, 0.75, 1.00], fontsize=7)
    # ax1.set_yticklabels([0, 0, 20, 40, 60], fontsize=7)
    # ax2 = fig.add_subplot(222)
    # ax2.set_title('Left Upper Part')
    # ax2.scatter(point_mre_x[200:], point_mare_y[200:], c=np.random.random(100), cmap='rainbow', s=16)
    # ax2.set_xlabel("MRE",fontsize=7)
    # ax2.set_ylabel("MARE",fontsize=7, labelpad=6)
    # ax2.set_xticklabels([0.00, 0.00, 0.50, 1.00], fontsize=7)
    # ax2.set_yticklabels([62.03, 62.03, 62.04, 62.05, 62.06], fontsize=7)
    # ax3 = fig.add_subplot(223)
    # ax3.set_title('Left Lower Part')
    # ax3.scatter(point_mre_x[100:200], point_mare_y[100:200], c=np.random.random(100), cmap='rainbow', s=16)
    # ax3.set_xlabel("MRE(1e-8+3.255e-4)",fontsize=7)
    # ax3.set_ylabel("MARE(1e-5+1.597e-1)",fontsize=7, labelpad=6)
    # ax3.set_xticklabels([3, 3, 4, 5], fontsize=7)
    # ax3.set_yticklabels([2, 2, 4, 6, 8], fontsize=7)
    # ax4 = fig.add_subplot(224)
    # ax4.set_title('Right Lower Part')
    # ax4.scatter(point_mre_x[:100], point_mare_y[:100], c=np.random.random(100), cmap='rainbow', s=16)
    # ax4.set_xlabel("MRE",fontsize=7)
    # ax4.set_ylabel("MARE",fontsize=7, labelpad=6)
    # ax4.set_xticklabels([0.950, 0.950, 0.975, 1.000, 1.025, 1.050], fontsize=7)
    # ax4.set_yticklabels([0.950, 0.950, 0.975, 1.000, 1.025, 1.050], fontsize=7)

    x_lst = []
    y_lst = []
    for i in range(len(point_mre_x)):
        if point_mare_y[i] < 1.5 and point_mre_x[i] < 0.2:
            x_lst.append(point_mre_x[i])
            y_lst.append(point_mare_y[i])

    print(max(point_mre_x))
    print(max(point_mare_y))
    print(min(point_mare_y))
    print(min(point_mare_y))

    # fig = plt.figure()
    # fig.suptitle('MRE/MARE Distribution Generated By Weighted Sampling On Batch_Norm')
    # fig.subplots_adjust(hspace=0.4)
    # ax1 = fig.add_subplot(121)
    # ax1.set_title('All Samples')
    # ax1.scatter(point_mre_x, point_mare_y, c=np.random.random(300), cmap='rainbow')
    # ax1.set_xlabel("MRE",fontsize=7)
    # ax1.set_ylabel("MARE",fontsize=7, labelpad=6)
    # ax1.set_xticklabels([0, 0, 250, 500, 750, 1000], fontsize=7)
    # ax1.set_yticklabels([0, 0, 5000, 10000, 15000, 20000, 25000], fontsize=7)
    # ax2 = fig.add_subplot(122)
    # ax2.set_title('Left Lower Part')
    # ax2.scatter(x_lst, y_lst, c=np.random.random(294), cmap='rainbow')
    # ax2.set_xlabel("MRE", fontsize=7)
    # ax2.set_ylabel("MARE", fontsize=7, labelpad=6)
    # ax2.set_xticklabels([0, 0, 100, 200, 300], fontsize=7)
    # ax2.set_yticklabels([0, 0, 200, 400, 600, 800, 1000], fontsize=7)

    fig = plt.figure()
    # fig.suptitle('MRE/MARE Distribution Generated By Weighted Sampling On Batch_Norm')
    fig.subplots_adjust(hspace=0.4)
    ax1 = fig.add_subplot(121)
    # ax1.set_title('All Samples')
    ax1.scatter(point_mre_x, point_mare_y, c=np.random.random(300), cmap='rainbow')
    ax1.set_xlabel("MRE",fontsize=7)
    ax1.set_ylabel("MARE",fontsize=7, labelpad=6)
    # ax1.set_xticklabels([0, 0, 250, 500, 750, 1000], fontsize=7)
    # ax1.set_yticklabels([0, 0, 5, 10, 15, 20, 25], fontsize=7)

    ax2 = fig.add_subplot(122)
    # ax2.set_title('Left Lower Part')
    ax2.scatter(x_lst, y_lst, c=np.random.random(199), cmap='rainbow')
    ax2.set_xlabel("MRE", fontsize=7)
    ax2.set_ylabel("MARE", fontsize=7, labelpad=6)
    # ax2.set_xticklabels([0, 0, 100, 200, 300], fontsize=7)
    # ax2.set_yticklabels([0, 0, 200, 400, 600, 800, 1000], fontsize=7)

    plt.tight_layout()
    plt.show()
    fig.savefig('./random.eps', format='eps', dpi=1000)
